package heap

// 堆
//
// 定义：
// - 堆是一个完全二叉树
// - 每一个堆节点都大于等于（或小于等于）其子节点
//
// 操作：
// - 插入一个元素 O(log(n))
// - 删除堆顶元素 O(log(n))
// - 排序（大顶堆、小顶堆） O(n * log(n))
//
// 堆的应用：
// - 优先级队列
// - 求Top K
// - 求中位数

type bigHeap struct {
	data []int
	// 堆节点数
	n int
	// 元素数
	count int
}

func NewHeap(size int) *bigHeap {
	data := make([]int, size+1)
	return &bigHeap{
		data:  data,
		n:     size,
		count: 0,
	}
}

type OutOfRangeError string

func (err OutOfRangeError) Error() string {
	return "The heap has no more space."
}

// 前序遍历
func (heap *bigHeap) ForEach(fn func(index int, ele int)) {
	for i := 1; i <= heap.n; i++ {
		fn(i, heap.data[i])
	}
}

// 插入: 插入，从下向上堆化
func (heap *bigHeap) Insert(val int) {
	if heap.count >= heap.n {
		panic(new(OutOfRangeError))
	}
	heap.count++
	heap.data[heap.count] = val
	heapify(heap.data, heap.count)
}

func heapify(data []int, index int) {
	for index/2 > 0 && data[index/2] < data[index] {
		swap(data, index/2, index)
		index = index / 2
	}
}

func swap(data []int, index1 int, index2 int) {
	tmp := data[index1]
	data[index1] = data[index2]
	data[index2] = tmp
}

// 删除堆顶元素 删除堆顶元素，然后将最后一个节点放到堆顶，从上向下堆化
func (heap *bigHeap) RemoveTop() {
	heap.data[1] = heap.data[heap.count]
	heap.data[heap.count] = -1
	heap.count--
	heapifyFromTop(heap.data, heap.n, 1)
}

// 排序：
// - 构建堆：插入思维；或从n/2到1个节点进行从上向下堆化
// - 排序：将堆顶元素放到最后一位，将最后一位元素放到堆顶，对1到n-1元素进行从上向下堆化。
func (heap *bigHeap) Sort() {
	if heap.count < 2 {
		return
	}
	buildHeap(heap.data, heap.n)
	for i := heap.n; i > 1; {
		swap(heap.data, 1, i)
		i--
		heapifyFromTop(heap.data, i, 1)
	}
}

func buildHeap(data []int, n int) {
	for start := n / 2; start > 0; start-- {
		index := start
		for {
			maxPos := index
			if 2*index <= n && data[maxPos] < data[2*index] {
				maxPos = 2 * index
			}
			if 2*index+1 <= n && data[maxPos] < data[2*index+1] {
				maxPos = 2*index + 1
			}
			if maxPos == index {
				break
			}
			swap(data, maxPos, index)
			index = maxPos
		}

	}
}

func heapifyFromTop(data []int, indexRange int, startCount int) {
	index := startCount
	for {
		maxPos := index
		if 2*index <= indexRange && data[maxPos] < data[2*index] {
			maxPos = 2 * index
		}
		if 2*index+1 <= indexRange && data[maxPos] < data[2*index+1] {
			maxPos = 2*index + 1
		}
		if maxPos == index {
			break
		}
		swap(data, maxPos, index)
		index = maxPos
	}
}
